#!/usr/bin/env python
# encoding: utf-8
"""
@version : 
@author  : 
@license : 
@contact : ****@massclouds.com
@site    : http://blog.csdn.net/***
@software: PyCharm
@time    : 17-2-13 上午10:33
"""

'''
apt-get install python-libvirt
apt-get install python-lxml
'''

import os
import logging

import libvirt
from lxml import etree
import subprocess

from kvmgt_template import KVMGT_TEMPLATE_XML_STR

class VirtMgr(object):
    '''
    classdocs
    '''
    # uri='qemu://root@192.168.100.219/system'
    def __init__(self, uri='qemu:///system'):
        ''' Constructor '''
        self.__conn = None
        self._uri = uri

    @property
    def conn(self):
        if not self.__conn:
            self.__conn = libvirt.open(self._uri)
        return self.__conn

    def listAllDomains(self):
        return self.conn.listAllDomains()

    def listDomStatus(self):
        vmstatus = {}
        for vm in self.conn.listAllDomains():
            vmstatus[vm.name()] = False
            if vm.isActive():
                vmstatus[vm.name()] = True
        return vmstatus

    def domainCnt(self):
        return len(self.conn.listAllDomains())

    def getDomByName(self, name):
        for dm in self.conn.listAllDomains():
            if dm.name() == name :
                return dm
        return None

    def getDomByUuid(self,uuid):
        for dm in self.conn.listAllDomains():
            if dm.UUIDString() == uuid:
                return dm
        return None


    #
    # def getMemByDomain(self,domain):
    #     res = domain.memoryStats()
    #     #return (res['actual']/1024)
    #     return res['actual']

    def getMaxMemByDomain(self,domain):
        xmlstr = domain.XMLDesc()
        dom = etree.XML(xmlstr)
        for disk in dom.xpath('/domain/memory'):
            # print etree.SubElement(disk, 'memory')
            # print ": ",etree.tostring(disk)
            # print disk.tag
            # print disk.text
            # print disk.get('unit') # print disk.values()
            return disk.text + " " + disk.get('unit')

    def getMemByDomain(self,domain):
        # return domain.maxMemory()
        xmlstr = domain.XMLDesc()
        dom = etree.XML(xmlstr)
        for disk in dom.xpath('/domain/currentMemory'):
            return disk.text + " " + disk.get('unit')

    def getCpusByDomain(self,domain):
        xmlstr = domain.XMLDesc()
        dom = etree.XML(xmlstr)
        for disk in dom.xpath('/domain/vcpu'):
            return disk.text



    # def getVcpus(self,domain):
    #     # ([(0, 1, 224810000000L, 2)], [(True, True, True, True)])
    #     return domain.vcpus()[0][0]
    #
    # def getmaxVcpus(self,domain):
    #     return domain.maxVcpus()


    def getDiskByDomain(self,domain):
        disks = []
        xmlstr = domain.XMLDesc()
        dom = etree.XML(xmlstr)
        for disk in dom.xpath('/domain/devices/disk'):
            if disk.get('type') == 'file' and \
                disk.get('device', 'disk') in ['disk', 'cdrom', 'floppy']:
                disk_path = disk.xpath('./source/@file')[0]
                # print os.path.getsize(disk_path) / 1024
                # os.system("qemu-img info %s --output=json" %disk_path)
                # cmd = [ "qemu-img", "info", disk_path,"--output=json"]
                # pro = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                # print pro.stdout.readall()
                # import json
                # print json.loads(pro.stdout.readlines())
                # print pro.stdout.readlines()," \n"


                disks.append((disk.get('device', 'disk'), disk_path))
        return disks

    def get_template_xml(self):
        # TPL_XML_PATH = '/opt/massclouds/share/template/template.xml'

        # return KVMGT_TEMPLATE_XML_STR

        TPL_XML_PATH = '/cdata/template/kvmgt-template.xml'
        if not os.path.exists(TPL_XML_PATH):
            logging.error("kvmgt-template.xml: %s not found, use default cnfig")
            # raise Exception('kvmgt-template.xml: %s not found')
            return KVMGT_TEMPLATE_XML_STR

        with open(TPL_XML_PATH) as fp:
            return fp.read()

    # def start_vm(self, vm_name):
    #     try:
    #         vm_domain = self.get_vm(vm_name)
    #         vm_domain.create()
    #     except Exception as e:
    #         logging.exception(e)
    #         raise Exception(e)

    # def shutdown_vm(self, vm_name):
    #     try:
    #         vm = self.get_vm(vm_name)
    #         if vm.isActive():
    #             vm.shutdown()
    #     except Exception as e:
    #         logging.exception(e)
    #
    # def poweroff_vm(self, vm_name):
    #     try:
    #         vm = self.get_vm(vm_name)
    #         if vm.isActive():
    #             vm.destroy()
    #     except Exception as e:
    #         logging.exception(e)
    #         raise Exception(e)

    def create_vm(self, **kwargs):
        '''
        kwargs = {
                name:None,           # set vm name
                vcpus:1,             # set num of vcpu
                memory:1024,         # set mem (unit: MiB)
                cdroms:[],           # set cdroms (path to iso files)
                disks:[],            # set disks (path to qcow2 files)
                disksize:50G,
                bridge:'br0',        # which bridge vm on
                xmlstr:None          # override all above (costom vm)
            }
        '''
        if not kwargs.get("name"):
            raise Exception("missing name information")

        def get_devno(prefix='sd'):
            devno = 'abcdefghijklmnopqrstuvwxyz'
            for newno in iter(devno):
                yield prefix + newno
        #
        # if kwargs.has_key('xmlstr'):
        #     self.conn.defineXML(kwargs('xmlstr'))
        #     return True

        xmlstr = self.get_template_xml()
        root = etree.XML(xmlstr)

        # Set vm name
        vmname = root.xpath('/domain/name')[0]
        vmname.text = kwargs['name']

        # Set vcpu
        vcpu = root.xpath('/domain/vcpu')[0]
        vcpu.text = str(kwargs.get('vcpus', 1))
        cpu = root.xpath('/domain/cpu')[0]
        topology = cpu.find('topology')
        topology.set("sockets", str(1))
        topology.set("cores", str(1))
        topology.set("threads", str(kwargs.get('vcpus', 1)))

        # Set memory
        mem = root.xpath('/domain/memory')[0]
        cmem = root.xpath('/domain/currentMemory')[0]
        mem.text = cmem.text = str(kwargs.get('memory', 1024))

        # Set DISK
        devices = root.xpath('/domain/devices')[0]

        dev = get_devno('sd')
        for image in kwargs.get('disks', []):
            disk = etree.Element('disk', {'type':'file', 'device':'disk'})
            disk.append(etree.Element('driver', {'name':'qemu', 'type':'qcow2'}))
            disk.append(etree.Element('source', {'file':os.path.abspath(image)}))
            disk.append(etree.Element('target', {'dev':dev.next(), 'bus':'virtio'}))
            devices.append(disk)

            # create image
            disksize = kwargs.get('disksize', "50G")
            cmd = ['qemu-img', 'create', '-f', 'qcow2',   os.path.abspath(image), disksize]
            p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            out, err = p.communicate()
            if p.returncode != 0:
                logging.error("Failed: create image %s - %s"% (dev, err))
                raise Exception("Failed: create image %s" %dev)

        # Set CDROM
        dev = get_devno('hd')
        for iso in kwargs.get('cdroms', []):
            disk = etree.Element('disk', {'type':'file', 'device':'cdrom'})
            disk.append(etree.Element('driver', {'name':'qemu', 'type':'raw'}))
            disk.append(etree.Element('source', {'file':os.path.abspath(iso)}))
            disk.append(etree.Element('target', {'dev':dev.next(), 'bus':'ide'}))
            devices.append(disk)

        # Set BRIDGE
        ifaces = root.xpath('/domain/devices/interface')
        for iface in ifaces:
            source = iface.find('source')
            source.set('bridge', kwargs.get('bridge', 'br0'))

        # create vm
        xmlstr = etree.tostring(root)
        print xmlstr

        self.conn.defineXML(xmlstr)
        return True

    def _get_all_vm_disks(self, vm):
        '''
            get all disks of specfic vm
            reutrn [ ({device_type}, {disk_path}), ...]
            return None, if vm not found
        '''
        disks = []
        xmlstr = vm.XMLDesc()
        dom = etree.XML(xmlstr)
        for disk in dom.xpath('/domain/devices/disk'):
            if disk.get('type') == 'file' and \
                disk.get('device', 'disk') in ['disk', 'cdrom', 'floppy']:
                disk_path = disk.xpath('./source/@file')[0]
                disks.append((disk.get('device', 'disk'), disk_path))
        return disks

    def get_all_vm_disks(self, vm_name):
        ''' get all disks of vm include disk,cdrom,floppy device
        '''

        vm = self.get_vm(vm_name)
        if vm:
            return self._get_all_vm_disks(vm)
        return None

    def get_vm_display_port(self, vm_name, allow_protocol=['spice']):
        ''' get remote connection port of vm
            reutrn [ ({protocol}, {port}), ...]
            return None, if vm not found
        '''
        vm = self.get_vm(vm_name)
        dpy_list = []
        if vm:
            xmlstr = vm.XMLDesc()
            dom = etree.XML(xmlstr)
            for dpy in dom.xpath('/domain/devices/graphics'):
                if dpy.get('type') in allow_protocol:
                    dpy_list.append((dpy.get('type'), dpy.get('port')))
            return dpy_list

        return None

    def delete_vm(self, vm_name):
        ''' safely delete a vm
            all disks in /cdata/vms of vm will be deleted
        '''
        vm = self.get_vm(vm_name)
        if vm:
            if vm.isActive():
                vm.destroy()
            for type, disk in self._get_all_vm_disks(vm):
                if type == 'disk' and  \
                        disk.startswith('/cdata/vms'):
                    if os.path.exists(disk):
                        os.remove(disk)
            vm.undefine()
            return True
        return None

    # def is_active(self, vm_name):
    #     ''' Return if the vm is ACTIVE
    #     '''
    #     vm = self.get_vm(vm_name)
    #     if vm:
    #         return vm.isActive()
    #     else:
    #         return False
    """
    [
    'name',
    'UUIDString',
    'UUID',
    'isActive',
    'OSType',
    'XMLDesc',
    'create','reboot', 'reset', 'resume',
    'shutdown',
    'destroy',
    'undefine',
    'state',
 
    'ID',

    'setMaxMemory',
    'setMemory',
    'setMemoryFlags',
    'setMemoryParameters', 'setMemoryStatsPeriod',
    'maxMemory',
    'memoryParameters','memoryPeek','memoryStats',

    'maxVcpus','pinVcpu', 'pinVcpuFlags',
    'setVcpus', 'setVcpusFlags',
    'vcpuPinInfo',
    'vcpus', 'vcpusFlags'

    'listAllSnapshots',
    'shutdownFlags',
    'connect','info',
    'suspend',  'autostart', ' 'snapshotCurrent', 'snapshotListNames', 'snapshotLookupByName', 'snapshotNum',
        'undefineFlags', 'updateDeviceFlags',
      ]
    """

# /extdisk/iso/ubuntu-14.04.3-desktop-amd64.iso

def test_mian():
    kwargs ={
            "name"  :"test_win7_03",   # set vm name
            "vcpus" :1,             # set num of vcpu
            "memory":"1024",        # set mem (unit: MiB)
            "cdroms":[              # set cdroms (path to iso files)
                "/root/Windows-7-Ultimate-SP1-x64-677408.iso",
                "/root/virtio-win.iso",
                      ],
            "disks" :["/cdata/vms/test_win7_03.qcow2"],            # set disks (path to qcow2 files)

            # "cdroms":["/extdisk/iso/ubuntu-14.04.3-desktop-amd64.iso"],       # set cdroms (path to iso files)
            # "disks" :["/abcddd.qcow2"],            # set disks (path to qcow2 files)
            "disksize":"40G",
            "bridge":'br0',         # which bridge vm on
            "xmlstr":None           # override all above (costom vm)
            }
    VirtMgr().create_vm(**kwargs)


virt_mgr = VirtMgr()

def test_main2():
    #print VirtMgr().domSatus()
    #print VirtMgr().vm_cnt()
    domain = virt_mgr.listAllDomains()
    print virt_mgr.getDiskByDomain(domain[0])
    print "min mem :" ,virt_mgr.getMemByDomain(domain[0])
    print "max mem :" ,virt_mgr.getMaxMemByDomain(domain[0])
    print "cpus    :" ,virt_mgr.getCpusByDomain(domain[0])

if __name__ == '__main__':
    # test_mian()
    test_main2()


